mod domain;
mod utils;

use crate::domain::CONFIG;
use chrono::Utc;
use comfy_table::{Cell, CellAlignment, Table};
use csv::Reader;
use dashmap::DashMap;
use rayon::prelude::*;
use std::collections::HashMap;
use std::fs;
use std::fs::File;
use std::io::BufReader;
use std::path::Path;
use std::process::exit;
use std::sync::LazyLock;

/// (开仓价格, 持仓币数, 持仓方向, 持仓时间)
type Pos = (f64, f64, String, i64);

/// Key: Binance_副所_交易对  Value: (盈亏百分比, open_roe)
static RESULT: LazyLock<DashMap<String, (f64, f64)>> = LazyLock::new(|| DashMap::new());

fn list_files(path: &str) -> std::io::Result<Vec<(String, String, String, String)>> {
    let mut temp_list = Vec::with_capacity(4);

    for entry in fs::read_dir(Path::new(path))? {
        let entry = entry?;
        let path = entry.path();
        if path.is_file() {
            // 仅筛选文件，忽略目录

            let current_path = path.to_string_lossy().to_owned().to_string();

            // 文件绝对路径
            let file_name = path.file_name().unwrap().to_str().unwrap();

            // (主所, 副所, 交易对)
            let len = &file_name[..file_name.len() - 4]
                .split("_")
                .collect::<Vec<_>>();

            // (文件绝对路径, 主所, 副所, 交易对)
            temp_list.push((
                current_path,
                len[0].to_string(),
                len[1].to_string(),
                len[2].to_string(),
            ));
        }
    }
    Ok(temp_list)
}

fn get_profit(latest_price: f64, pos: &Pos) -> f64 {
    if pos.2 == "BUY" {
        (pos.1 * latest_price) - (pos.1 * pos.0)
    } else {
        (pos.1 * pos.0) - (pos.1 * latest_price)
    }
}

fn open(
    is_stop: &mut HashMap<String, i64>,
    next_open_timestamp: &mut Option<i64>,
    side: &str,
    coinex_timestamp: i64,
    is_open: bool,
    is_close: bool,
    coinex_bid: f64,
    coinex_ask: f64,
    balance: &mut f64,
    pos: &mut Option<Pos>,
    main_bid_amount: f64,
    main_ask_amount: f64,
    main_timestamp: i64,
    keep_time: &mut Option<i64>,
) {
    // 判断是否处于强停状态
    let stop_timestamp = is_stop.get(side);
    if stop_timestamp.is_some() {
        if coinex_timestamp < *stop_timestamp.unwrap() {
            return;
        }
    }

    // 未到下次开仓时间
    if next_open_timestamp.is_some() {
        if coinex_timestamp < next_open_timestamp.unwrap() {
            return;
        }
    }

    // 未达到开仓或平仓标准
    if !is_open || is_close {
        return;
    }

    // 做多情况下
    // 买单 必须大于等于 卖单 x倍
    if side == "BUY" && CONFIG.buy_multiple.is_some() {
        // 不满足
        if main_bid_amount / main_ask_amount < CONFIG.buy_multiple.unwrap() {
            return;
        }

        // 检查保持时间是否大于 x 毫秒
        if keep_time.is_some() {
            // 未达到 x 毫秒
            if (main_timestamp - keep_time.unwrap()) < CONFIG.keep_time.unwrap() {
                return;
            }

            // 满足 保持 x 毫秒, 重置时间
            *keep_time = None;
        }
        // 记录时间
        else {
            *keep_time = Some(main_timestamp);
            return;
        }
    }

    // 计算副所的平均价
    let coinex_avg_open_price = (coinex_bid + coinex_ask) / 2.0;
    // 计算下单量
    let amount = (*balance / coinex_avg_open_price).round();

    // 做多, 取副所卖1
    // 做空, 取副所买1
    let real_price = if side == "BUY" {
        coinex_ask
    } else {
        coinex_bid
    };

    // 更新余额
    *balance = *balance - (real_price * amount);

    // 模拟开仓 (开仓价格, 持仓币数, 持仓方向, 持仓时间)
    *pos = Some((real_price, amount, side.to_string(), coinex_timestamp));
    // println!("{:?}", pos);
}

fn close(
    pos: &mut Option<Pos>,
    coinex_bid: f64,
    coinex_ask: f64,
    main_bid: f64,
    main_ask: f64,
    coinex_timestamp: i64,
    is_stop: &mut HashMap<String, i64>,
    side: &str,
    is_close: bool,
    balance: &mut f64,
) {
    // 提取出当前的仓位 (开仓价格, 持仓币数, 持仓方向, 持仓时间)
    let current_pos = pos.as_ref().unwrap();

    // 不是同个方向的仓位, 跳过
    if side != current_pos.2 {
        return;
    }

    // 平仓价格
    let close_price: f64;

    // 计算实时的盈亏
    let current_profit = if current_pos.2 == "BUY" {
        close_price = coinex_bid;
        get_profit(main_bid, current_pos)
    } else {
        close_price = coinex_ask;
        get_profit(main_ask, current_pos)
    };

    // 计算当前仓位的持仓时间(毫秒)
    let current_hold_time = coinex_timestamp - current_pos.3;

    // 收益 / 实时仓位价值
    let real_value = coinex_bid * current_pos.1;
    let roe = current_profit / real_value;

    // 止损(止损时间 * 小时)
    if roe * 100.0 <= CONFIG.zhi_sun {
        // 设置止损时长
        is_stop.insert(
            side.to_string(),
            coinex_timestamp + (CONFIG.zhi_sun_time * 60000),
        );
        // println!("{side} {symbol} 触发止损");

        // 清空仓位, 计算余额
        *balance += real_value;
        *pos = None;

        return;
    }

    // 不允许平仓
    if !is_close {
        return;
    }

    /*
    (价格回归)
     */

    // 未达到最小持仓时间
    if current_hold_time < CONFIG.min_hold_time {
        return;
    }

    // 进行平仓
    *balance += current_pos.1 * close_price;
    *pos = None;
}

fn read_csv_stream(path: &str, coinex_exchange: &str, symbol: &str, open_roe: f64) {
    let reader = BufReader::new(File::open(path).unwrap());
    let mut rdr = Reader::from_reader(reader);

    /*
    [
      主所买1, 主所卖1, 主所买1挂单量, 主所卖1挂单量, 主所时间戳
      副所买1, 副所卖1, 副所买1挂单量, 副所卖1挂单量, 副所时间戳
    ]
    */

    // 模拟本金 1000 USDT
    let mut balance = 1000.0;
    // 仓位  (开仓价格, 持仓币数, 持仓方向, 持仓时间)
    let mut pos: Option<Pos> = None;
    // 强停时间
    let mut is_stop: HashMap<String, i64> = HashMap::new();
    // 下次开仓时间
    let mut next_open_timestamp: Option<i64> = None;
    let mut keep_time: Option<i64> = None;

    // 逐行迭代
    for result in rdr.records() {
        let Ok(record) = result else {
            println!("跳过");
            continue;
        };
        let main_bid: f64 = record.get(0).unwrap().parse().unwrap();
        let main_ask: f64 = record.get(1).unwrap().parse().unwrap();
        let main_bid_amount: f64 = record.get(2).unwrap().parse().unwrap();
        let main_ask_amount: f64 = record.get(3).unwrap().parse().unwrap();
        let main_timestamp: i64 = record.get(4).unwrap().parse().unwrap();
        let coinex_bid: f64 = record.get(5).unwrap().parse().unwrap();
        let coinex_ask: f64 = record.get(6).unwrap().parse().unwrap();
        // let coinex_bid_amount: f64 = record.get(7).unwrap().parse().unwrap();
        // let coinex_ask_amount: f64 = record.get(8).unwrap().parse().unwrap();
        let coinex_timestamp: i64 = record.get(9).unwrap().parse().unwrap();

        let buy = (main_bid - coinex_ask) / ((main_bid + coinex_ask) / 2.0) * 100.0;
        let sell = (main_ask - coinex_bid) / ((main_ask + coinex_bid) / 2.0) * 100.0 * -1.0;

        // println!("{buy} {sell}");

        for side in &CONFIG.side {
            // 是否允许开平仓
            let (is_open, is_close) = if side == "BUY" {
                // (buy >= open_roe, sell >= CONFIG.close_roe)
                (buy >= open_roe, sell >= 0.0)
            } else {
                // (sell >= open_roe, buy >= CONFIG.close_roe)
                (sell >= open_roe, buy >= 0.0)
            };

            // 准备开仓
            if pos.is_none() {
                open(
                    &mut is_stop,
                    &mut next_open_timestamp,
                    side,
                    coinex_timestamp,
                    is_open,
                    is_close,
                    coinex_bid,
                    coinex_ask,
                    &mut balance,
                    &mut pos,
                    main_bid_amount,
                    main_ask_amount,
                    main_timestamp,
                    &mut keep_time,
                );
            }
            // 准备平仓
            else {
                close(
                    &mut pos,
                    coinex_bid,
                    coinex_ask,
                    main_bid,
                    main_ask,
                    coinex_timestamp,
                    &mut is_stop,
                    side,
                    is_close,
                    &mut balance,
                );
            }
        }
    }

    // 计算盈亏百分比

    let final_profit = if balance > 1000.0 {
        ((balance - 1000.0) / 1000.0) * 100.0
    } else {
        (((1000.0 - balance) * -1.0) / 1000.0) * 100.0
    };

    let key = format!("Binance_{coinex_exchange}_{symbol}");

    let Some(mut result) = RESULT.get_mut(&key) else {
        RESULT.insert(key, (final_profit, open_roe));
        return;
    };

    if final_profit > result.0 {
        *result = (final_profit, open_roe);
    }
}

fn main() {
    let start_timestamp = Utc::now().timestamp();
    // 做空, 做多    0.1-0.2   =  2*10

    let Ok(task) = list_files(CONFIG.read_path.as_str()) else {
        println!("没有数据, 退出程序");
        exit(0);
    };

    // 0.10-0.20
    let range = (10..=20).map(|x| x as f64 / 100.0);

    let mut list = Vec::with_capacity(task.len() * 11);

    for item in task {
        for roe in range.to_owned() {
            list.push((roe, item.clone()));
        }
    }

    list.par_iter().for_each(|temp_task| {
        let file = Path::new(&CONFIG.read_path)
            .join(temp_task.1.0.as_str())
            .to_str()
            .unwrap()
            .to_string();

        read_csv_stream(
            file.as_str(),
            temp_task.1.2.as_str(),
            temp_task.1.3.as_str(),
            temp_task.0,
        )
    });

    let end_timestamp = Utc::now().timestamp() - start_timestamp;

    // (主所, 副所, 交易对, 盈亏百分比, open_roe)
    let mut temp_list = Vec::with_capacity(RESULT.len());

    for item in RESULT.iter() {
        let mut parts = item.key().split("_");
        let main_exchange = parts.next().unwrap();
        let coinex_exchange = parts.next().unwrap();
        let symbol = parts.next().unwrap();

        temp_list.push((
            main_exchange.to_string(),
            coinex_exchange.to_string(),
            symbol.to_string(),
            item.value().0,
            item.value().1,
        ));
    }

    temp_list.sort_by(|a, b| a.3.partial_cmp(&b.3).unwrap());
    temp_list.reverse();

    let mut table = Table::new();
    table.set_header(vec!["main", "coinex", "symbol", "profit", "roe"]);

    // 打印排序后的结果
    for item in temp_list.iter() {
        table.add_row(vec![
            Cell::new(&item.0).set_alignment(CellAlignment::Center),
            Cell::new(&item.1).set_alignment(CellAlignment::Center),
            Cell::new(&item.2).set_alignment(CellAlignment::Center),
            Cell::new(format!("{:.2}", item.3)).set_alignment(CellAlignment::Center),
            Cell::new(item.4).set_alignment(CellAlignment::Center),
        ]);
    }

    println!("{table}");
    println!("回测耗时: {end_timestamp} 秒")
}
